from strange_io import SIGNATURE, _load_bytearray, get_format
from strange_io.colors import StrangeColors
from strange_io.array import StrangeArray
import strange_io.dynamic_integer as dynamic_integer

class StrangePixels:
    def __init__(self, width=20, height=20, palette_size:int=1):
        self.palette_size:int = palette_size
        self.height:int = height
        self.width:int = width
        self.image:list[list[int]] = [[0 for _ in range(height)] for _ in range(width)]

    def dump(self) -> dict:
        out = {}
        out['palette_size'] = self.palette_size
        out['height'] = self.height
        out['width'] = self.width
        out['image'] = self.image
        return out

    def load(self, image) -> dict:
        self.palette_size = image['palette_size']
        self.height = image['height']
        self.width = image['width']
        self.image = image['image']

    def dump(self) -> dict:
        out = {}
        out['palette_size'] = self.palette_size
        out['height'] = self.height
        out['width'] = self.width
        out['image'] = self.image
        return out

    def load(self, image) -> dict:
        self.palette_size = image['palette_size']
        self.height = image['height']
        self.width = image['width']
        self.image = image['image']

    def dump_bytes(self) -> bytes:
        return self.encode(self.dump())

    def load_bytes(self, image:bytes):
        self.load(self.decode(image))

    def dump_file(self, path):
        """ Dump document to file """
        out = bytearray()
        header = bytearray(SIGNATURE + b'pixels\n')
        out += header
        out += self.dump_bytes()
        with open(path, 'wb') as f:
            f.write(out)

    def load_file(self, path, ignore_type=False) -> dict:
        """ Load document from file """
        if get_format(path) == b'pixels' or ignore_type:
            file = _load_bytearray(path, remove_type=True)
            self.load_bytes(file)
        else:
            raise TypeError('Attempted to load non-pixels document as pixels document')

    @staticmethod
    def encode(img:dict) -> bytes:
        out = bytearray()
        out += img['palette_size'].to_bytes(1, 'big')
        out += dynamic_integer.encode(img['width'])[0]
        out += dynamic_integer.encode(img['height'])[0]
        out_image = []
        for x, column in enumerate(img['image']):
            column_data = b''
            for y, pixel in enumerate(column):
                column_data += pixel.to_bytes(1, 'big')
            out_image.append(column_data)
        sarray = StrangeArray(out_image)
        out += sarray.dump_bytes()
        return out

    @staticmethod
    def decode(file:bytes) -> dict:
        out = {}
        out['palette_size'] = file[0]
        del file[0]
        width = dynamic_integer.decode(file)[0]
        out['width'] = width
        height = dynamic_integer.decode(file)[0]
        out['height'] = height
        out['image'] = [[0 for _ in range(height)] for _ in range(width)]
        column = []
        sarray = StrangeArray()
        sarray.load_bytes(file)
        for x, column in enumerate(sarray.dump()):
            for y, pixel in enumerate(column):
                out['image'][x][y] = pixel
        return out
